#include "vector_database.h"
#include "index_factory.h"
#include "hnswlib_index.h"
#include "faiss_index.h"
#include "index_factory.h"
VectorDatabase::VectorDatabase(const std::string& db_path,const std::string &wal_path):scalar_storage_(db_path){
    persistence_.init(wal_path);
}
// 插入或更新向量
void VectorDatabase::upsert_vectors(uint64_t id,const rapidjson::Document& data,IndexFactory::IndexType index_type){
    // 检查标量存储中是否有该id
    rapidjson::Document existData;
    try{
        existData=scalar_storage_.get_scalar(id);
    }catch(std::exception& e){
        //todo: 如果没有，则插入
    }

    // 如果存在，则从索引删除
    if(existData.IsObject()){
        std::vector<float> existVector(existData["vector"].Size());
        for(rapidjson::SizeType i=0;i<existVector.size();i++){
            existVector[i]=existData["vector"][i].GetFloat();
        }
        void* index=IndexFactory::getGlobalIndexFactory().getIndex(index_type);
        switch(index_type){
            case IndexFactory::IndexType::HNSW:
                {HNSWLibIndex* hnswLibIndex=static_cast<HNSWLibIndex*>(index);
                // todo: 删除向量
                break;}
            case IndexFactory::IndexType::FLAT:
                {FaissIndex* faissIndex=static_cast<FaissIndex*>(index);
                faissIndex->remove_vectors({static_cast<long>(id)});
                break;}
            default:
                break;
        }
    }

    // 新向量插入索引
    std::vector<float> newVector(data["vectors"].Size());
    for(rapidjson::SizeType i=0;i<data["vectors"].Size();i++){
        newVector[i]=data["vectors"][i].GetFloat();
    }

    void* index=IndexFactory::getGlobalIndexFactory().getIndex(index_type);
    switch (index_type)
    {
    case IndexFactory::IndexType::HNSW:
        {HNSWLibIndex* hnswIndex=static_cast<HNSWLibIndex*>(index);
        hnswIndex->insert_vectors(newVector,id);
        break;}
    
    case IndexFactory::IndexType::FLAT:
        {FaissIndex* flatIndex=static_cast<FaissIndex*>(index);
        flatIndex->insert(newVector,id);
        break;}
    default:
        break;
    }
    // 更新标量存储中的向量
    scalar_storage_.insert_scalar(id,data);
}
rapidjson::Document VectorDatabase::query(uint64_t id){
    return scalar_storage_.get_scalar(id);
}

std::pair<std::vector<long>,std::vector<float>> VectorDatabase::search(const rapidjson::Document& json_request){
    // 解析请求参数
    std::vector<float> query;
    for(const auto&it:json_request[REQUEST_VECTORS].GetArray()){
        query.push_back(it.GetFloat());
    }

    int k=json_request[REQUEST_K].GetInt();

    //获取索引类型
    IndexFactory::IndexType index_type=IndexFactory::IndexType::UNKNOWN;
    if(json_request.HasMember(REQUEST_INDEX_TYPE)&&json_request[REQUEST_INDEX_TYPE].IsString()){
        std::string index_type_str=json_request[REQUEST_INDEX_TYPE].GetString();
        if(index_type_str==INDEX_TYPE_FLAT){
            index_type=IndexFactory::IndexType::FLAT;
        }else if(index_type_str==INDEX_TYPE_HNSW){
            index_type=IndexFactory::IndexType::HNSW;
        }
    }

    //检查是否包含filter参数
    roaring_bitmap_t* filter_bitmap=nullptr;
    if(json_request.HasMember("filter")&&json_request["filter"].IsObject()){
        const auto& filter=json_request["filter"];
        std::string fieldName=filter["fieldName"].GetString();
        std::string op_str=filter["op"].GetString();
        int64_t value=filter["value"].GetInt64();

        FilterIndex::Operation op=(op_str=="=")?FilterIndex::Operation::EQUAL:FilterIndex::Operation::NOT_EQUAL;

        //获取FilterIndex
        FilterIndex* filter_index=static_cast<FilterIndex*>(IndexFactory::getGlobalIndexFactory().getIndex(IndexFactory::IndexType::FILTER));
        filter_bitmap=roaring_bitmap_create();
        filter_index->getIntFieldFilter(fieldName,value,op,filter_bitmap);
    }

    void* index=IndexFactory::getGlobalIndexFactory().getIndex(index_type);

    std::pair<std::vector<long>,std::vector<float>> results;
    switch (index_type)
    {
    case IndexFactory::IndexType::FLAT:
        {FaissIndex* faiss_index=static_cast<FaissIndex*>(index);
        results=faiss_index->searchVector(query,k,filter_bitmap);
        break;}
    case IndexFactory::IndexType::HNSW:{
        HNSWLibIndex* hnsw_index=static_cast<HNSWLibIndex*>(index);
        results=hnsw_index->search_vectors(query,k);
        break;
    }
    default:
        break;
    }
    if(filter_bitmap!=nullptr){
        delete filter_bitmap;
    }
    return results;
}

void VectorDatabase::reloadDatabase(){
    GlobalLogger->debug("Reloading database...");
    std::string operation_type;
    rapidjson::Document json_data;
    persistence_.readNextWALLog(&operation_type,&json_data);

    while(!operation_type.empty()){
        GlobalLogger->debug("Reloading database operation: {}",operation_type);

        // 打印读取内容
        rapidjson::StringBuffer buffer;
        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
        json_data.Accept(writer);
        GlobalLogger->debug("Reloading database content: {}",buffer.GetString());
        // 只关注upsert
        if(operation_type=="upsert"){
            uint64_t id=json_data[REQUEST_ID].GetUint64();
            IndexFactory::IndexType index_type=getIndexTypeFromReq(json_data);
            upsert_vectors(id,json_data,index_type);
        }

        rapidjson::Document().Swap(json_data);
        operation_type.clear();
        persistence_.readNextWALLog(&operation_type,&json_data);
    }
}
void VectorDatabase::writeWALLog(const std::string& operation_type,const rapidjson::Document& json_data){
    std::string version="1.0.0";
    persistence_.writeWALLog(operation_type,json_data,version);
}

void VectorDatabase::writeWALLogWithID(uint64_t id, const std::string &data)
{
    std::string operation_type="upsert";
    std::string version="1.0.0";
    persistence_.writeWALRawLog(id,operation_type,data,version);
}

void VectorDatabase::takeSnapshot()
{
    persistence_.takeSnapshot(scalar_storage_);
}

int64_t VectorDatabase::getStartIndexID() const
{
    return persistence_.getID();
}

IndexFactory::IndexType VectorDatabase::getIndexTypeFromReq(const rapidjson::Document& json_request){
    if(!json_request.HasMember(REQUEST_INDEX_TYPE))return IndexFactory::IndexType::UNKNOWN;
    std::string index_type=json_request[REQUEST_INDEX_TYPE].GetString();
    if(index_type==INDEX_TYPE_FLAT){
        return IndexFactory::IndexType::FLAT;
    }else if(index_type==INDEX_TYPE_HNSW){
        return IndexFactory::IndexType::HNSW;
    }
}